from math import log
def createDataSet():
    dataSet = [[1,1,'yes'],
               [1,1,'yes'],
               [1,0, 'no'],
               [0,1,'no'],
               [0,1,'no']]
    labels = ['no serfacing', 'flippers']
    return dataSet, labels
"""
    计算香农熵
    :return 香农熵
"""
def calcShannonEnt(dataSet):
    numEntries = len(dataSet)
    countLabels = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        #print(currentLabel)
        if currentLabel not in countLabels.keys():
            countLabels[currentLabel] = 0
        countLabels[currentLabel] += 1
    shannonEnt = 0.0
    #print(countLabels)
    for key in countLabels:
        prop = float(countLabels[key]) / numEntries
        shannonEnt -= prop * log(prop, 2)
    return shannonEnt
"""
    参数：数据集、划分数据集特征、返回特征值
    :return 特征值
"""
def splitDataSet(dataSet, axis, value):
    retList = []
    for featVec in dataSet:
        if featVec[axis] == value:
            reduceFeatVec = featVec[:axis]
            reduceFeatVec.extend(featVec[axis+1:])
            retList.append(reduceFeatVec)
    return retList
"""
选取最佳特征用来分类
return 最佳分类特征的下标
"""
def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0]) -  1
    baseEntropy = calcShannonEnt(dataSet)
    bestInfoGain = 0.0
    bestFeature = -1
    for i in range(numFeatures):
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)#创建唯一分类标签列表
        #print (uniqueVals)
        newEntropy = 0.0
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet, i, value)
            #print(subDataSet)
            prob = len(subDataSet) / len(dataSet)
            newEntropy += prob * calcShannonEnt(subDataSet)
        infoGain = baseEntropy - newEntropy
        if (infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature
import operator
"""
    数据集处理完所有后，但类标签不唯一通过多数表决的方法定义叶子节点的分类
"""
def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys(): classCount[vote] = 0
        classCount[vote] += 1  #reverse=True表示按降序
    sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)#sorted返回一个新的list 不是在原有的基础上做修改
    return sortedClassCount[0][0]
"""
    递归创建决策树
    :return 决策树
"""
def createTree(dataSet, labels):
    classList = [example[-1] for example in dataSet]
    #print('ClassList: ',classList)
    if classList.count(classList[0]) == len(classList):#类别相同停止划分
        return classList[0]
    if len(dataSet[0]) == 1:
        #print('ClassList: ', classList)
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet)
    bestFeatLabel = labels[bestFeat]
    myTree = {bestFeatLabel:{}}
    del(labels[bestFeat])
    featValues = [example[bestFeat] for example in dataSet]
    #print('featValues: ', featValues)
    uniqueVals = set(featValues)
    for value in uniqueVals:
        subLabels = labels[:]
        #print('SubLabels: ', subLabels)
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
    return myTree
def classify(inputTree, featLabels, testVec):
    firstStr = list(inputTree.keys())[0]#python2.x中dict.keys返回一个列表。3.x中返回dict_keys对象
    secondDict = inputTree[firstStr]
    print(secondDict)
    featIndex = featLabels.index(firstStr)
    for key in list(secondDict.keys()):
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:
                classLabel = secondDict[key]
    return classLabel

"""
    绘制树形图
"""
import matplotlib.pyplot as plt
decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4",fc="0.8")
arrow_args = dict(arrowstyle='<-')
def plotNode(nodeText, centerPt, parentPt, nodeType):
    pass

"""
    决策树存储
"""
def storeTree(inputTree, filename):
    import pickle
    fw = open(filename, 'wb+')
    pickle.dump(inputTree, fw)
    fw.close()
"""
    读取决策树
    :return 决策树dict
"""
def grabTree(filename):
    import pickle
    fw = open(filename, 'rb')
    return pickle.load(fw)
dataSet, labels = createDataSet()
#calcShannonEnt(dataSet)
#print(calcShannonEnt(dataSet))
#dataSet[0][-1] = 'maybe'
#print(dataSet)
#print(calcShannonEnt(dataSet))
#print(splitDataSet(dataSet, 0, 0))
#print(chooseBestFeatureToSplit(dataSet))
#print (createTree(dataSet, labels))
inputTree = createTree(dataSet, labels)
print(inputTree)
#list = [('b', 2), ('a', 3), ('c', 1), ('d', 5)]
#print(sorted(list, key=lambda x:x[0], reverse=True))
#classlist = ['no', 'yes', 'no', 'yes', 'no']
#print(majorityCnt(classlist))
dataSet, labels = createDataSet()
print(classify(inputTree, labels, [1,1]))
storeTree(inputTree, 'classifierStorage.txt')
print(grabTree('classifierStorage.txt'))